package com.ksc.wordcount.driver;

import com.ksc.wordcount.rpc.ExecutorRegister;
import com.ksc.wordcount.task.TaskStatus;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class ExecutorManager {

    /**
     * ExecutorUrl和ExecutorRegister的映射
     */
    private Map<String, ExecutorRegister> executorRegisterMap = new HashMap<>();

    /**
     * ExecutorUrl和Core数的映射
     */
    private Map<String, Integer> executorAvailableCoresMap = new HashMap<>();

    /**
     * 更新Executor的注册信息
     *
     * @param executorRegister Executor的注册信息
     */
    public void updateExecutorRegister(ExecutorRegister executorRegister) {
        executorRegisterMap.put(executorRegister.getExecutorUrl(), executorRegister);
        // 建立ExecutorUrl和Core数的映射
        executorAvailableCoresMap.put(executorRegister.getExecutorUrl(), executorRegister.getCores());
    }

    /**
     * 获取Executor的可用资源核心数Map
     *
     * @return Executor的可用资源核心数Map
     */
    public Map<String, Integer> getExecutorAvailableCoresMap() {
        return executorAvailableCoresMap;
    }

    /**
     * 获取Executor的注册信息Map
     *
     * @return Executor的注册信息Map
     */
    public Map<String, ExecutorRegister> getExecutorRegisterMap() {
        return executorRegisterMap;
    }

    /**
     * 获取Executor的最大资源核心数
     *
     * @param executorUrl Executor的URL
     * @return Executor的最大资源核心数
     */
    public int getExecutorMaxCore(String executorUrl) {
        return executorRegisterMap.get(executorUrl).getCores();
    }

    /**
     * 获取Executor的可用资源核心数
     *
     * @param executorUrl Executor的URL
     * @return Executor的可用资源核心数
     */
    public int getExecutorAvalibeCore(String executorUrl) {
        return executorAvailableCoresMap.get(executorUrl);
    }

    /**
     * 更新Executor的可用资源核心数
     *
     * @param executorUrl Executor的URL
     * @param addCores    新增的核心数
     */
    public synchronized void updateExecutorAvailableCores(String executorUrl, int addCores) {
        int oldCore = executorAvailableCoresMap.get(executorUrl) == null ? 0 : executorAvailableCoresMap.get(executorUrl);
        executorAvailableCoresMap.put(executorUrl, oldCore + addCores);
    }
}